import os
import sys
from queue import Queue
import threading
import time
import tkinter as tk
from tkinter import filedialog, simpledialog, scrolledtext, messagebox, ttk
import json
import cv2
import numpy as np
from recognition_UI import Recognition
from midiMapping import createMidiInfo
from createmidi import createMidi
from dataset_generator import DatasetGenerator  # 确保导入DatasetGenerator

# 全局参数设置
args = {
    "speed": 150, "tonic": 'C', 'beat': "4/4", "delta_range": 1.5, "print_info": 1, "match_precision": 2,
    "thread_num": 8, "sampling_rate": 2, "threshold_method": "global", "threshold": 160
}

class MainApplication(tk.Tk):
    def __init__(self):
        super().__init__()
        self.sheet_recognition = None
        self.output_queue = Queue()
        self.title("简谱识别转MIDI")
        
        # 设置窗口大小和位置
        self.set_window_geometry()
        
        # 创建主界面控件
        self.create_widgets()
        
        # 重定向输出
        self.redirect_output()
    
    def set_window_geometry(self):
        """设置窗口大小和位置"""
        screen_width = self.winfo_screenwidth()
        screen_height = self.winfo_screenheight()
        window_width = 900  # 增加宽度以适应新控件
        window_height = 800  # 增加高度以确保所有控件可见
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2
        self.geometry(f"{window_width}x{window_height}+{x}+{y}")
        self.resizable(False, False)
    
    def create_widgets(self):
        """创建界面控件"""
        # 主框架
        main_frame = tk.Frame(self)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 文件选择部分
        self.create_file_selection(main_frame)
        
        # 参数设置部分
        self.create_parameter_settings(main_frame)
        
        # 输出区域
        self.create_output_area(main_frame)
        
        # 按钮区域 - 确保在输出区域下方
        self.create_button_area(main_frame)
    
    def create_file_selection(self, parent):
        """创建文件选择部分"""
        file_frame = tk.LabelFrame(parent, text="文件选择")
        file_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 路径标签和输入框
        tk.Label(file_frame, text="文件/文件夹路径:").grid(row=0, column=0, sticky="w", padx=5, pady=5)
        self.entry_path_var = tk.StringVar()
        self.entry_path = tk.Entry(file_frame, textvariable=self.entry_path_var, width=50)
        self.entry_path.grid(row=0, column=1, padx=5, pady=5, sticky="ew")
        
        # 浏览按钮
        self.button_browse = tk.Button(file_frame, text="浏览", command=self.browse_path)
        self.button_browse.grid(row=0, column=2, padx=5, pady=5)
    
    def create_parameter_settings(self, parent):
        """创建参数设置部分"""
        # 参数框架
        param_frame = tk.LabelFrame(parent, text="识别参数设置")
        param_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 参数设置 - 使用网格布局
        self.create_parameter_grid(param_frame)
        
        # 高级选项框架
        adv_frame = tk.LabelFrame(parent, text="高级选项")
        adv_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 高级选项 - 使用网格布局
        self.create_advanced_grid(adv_frame)
    
    def create_parameter_grid(self, frame):
        """创建主要参数网格"""
        row = 0
        
        # 主音
        tk.Label(frame, text="简谱主音(1=):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_tonic_var = tk.StringVar(value=args["tonic"])
        tk.Entry(frame, textvariable=self.entry_tonic_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 速度
        tk.Label(frame, text="简谱速度(每分钟拍数):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_speed_var = tk.StringVar(value=str(args["speed"]))
        tk.Entry(frame, textvariable=self.entry_speed_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 节拍
        tk.Label(frame, text="节拍(如4/4):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_beat_var = tk.StringVar(value=args["beat"])
        tk.Entry(frame, textvariable=self.entry_beat_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 输出中间信息
        tk.Label(frame, text="输出中间信息:").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.var_print_info = tk.BooleanVar(value=bool(args["print_info"]))
        tk.Checkbutton(frame, variable=self.var_print_info).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 匹配精度
        tk.Label(frame, text="匹配精度:").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.match_precision_var = tk.StringVar()
        self.match_precision_combo = ttk.Combobox(frame, textvariable=self.match_precision_var, 
                                                values=["高精度 (慢)", "中精度 (推荐)", "低精度 (快)"],
                                                width=15, state="readonly")
        self.match_precision_combo.current(1)  # 默认选择中精度
        self.match_precision_combo.grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
    
    def create_advanced_grid(self, frame):
        """创建高级选项网格"""
        row = 0
        
        # 线程数
        tk.Label(frame, text="线程数(1-16):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_thread_num_var = tk.StringVar(value=str(args["thread_num"]))
        tk.Entry(frame, textvariable=self.entry_thread_num_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 字符间距
        tk.Label(frame, text="字符上下间距(1.0-2.5):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_delta_range_var = tk.StringVar(value=str(args["delta_range"]))
        tk.Entry(frame, textvariable=self.entry_delta_range_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 二值化方法
        tk.Label(frame, text="二值化方法:").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.threshold_method_var = tk.StringVar()
        self.threshold_method_combo = ttk.Combobox(frame, textvariable=self.threshold_method_var, 
                                                  values=["全局阈值", "OTSU自适应"],
                                                  width=15, state="readonly")
        self.threshold_method_combo.current(0)  # 默认全局阈值
        self.threshold_method_combo.grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 全局阈值
        tk.Label(frame, text="全局阈值(0-255):").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.entry_threshold_var = tk.StringVar(value=str(args["threshold"]))
        tk.Entry(frame, textvariable=self.entry_threshold_var, width=10).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 手动调整
        tk.Label(frame, text="手动调整中间过程:").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.var_wanna_debug = tk.BooleanVar(value=False)
        tk.Checkbutton(frame, variable=self.var_wanna_debug).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
        
        # 启动标注工具
        tk.Label(frame, text="完成后启动标注工具:").grid(row=row, column=0, sticky="w", padx=5, pady=2)
        self.var_launch_annotation = tk.BooleanVar(value=True)
        tk.Checkbutton(frame, variable=self.var_launch_annotation).grid(row=row, column=1, padx=5, pady=2, sticky="w")
        row += 1
    
    def create_output_area(self, parent):
        """创建输出区域"""
        # 输出标签
        output_label = tk.Label(parent, text="程序输出:")
        output_label.pack(anchor="w", padx=10, pady=(10, 0))
        
        # 输出文本框
        self.output_text = scrolledtext.ScrolledText(parent, width=100, height=15, wrap=tk.WORD)  # 减少高度
        self.output_text.pack(fill=tk.BOTH, expand=True, padx=10, pady=(0, 5))
    
    def create_button_area(self, parent):
        """创建按钮区域"""
        # 创建按钮框架
        button_frame = tk.Frame(parent)
        button_frame.pack(fill=tk.X, padx=10, pady=10)
        
        # 运行按钮
        self.run_button = tk.Button(button_frame, text="开始识别", command=self.run_recognition, 
                                  bg="#4CAF50", fg="white", height=2, width=20)
        self.run_button.pack(side=tk.LEFT, padx=10)
        
        # 添加分隔空间
        tk.Label(button_frame, text="", width=5).pack(side=tk.LEFT)
        
        # 添加退出按钮
        exit_button = tk.Button(button_frame, text="退出程序", command=self.destroy, 
                             bg="#f44336", fg="white", height=2, width=20)
        exit_button.pack(side=tk.LEFT, padx=10)
    
    def redirect_output(self):
        """重定向输出到文本框"""
        sys.stdout = TextRedirector(self.output_text, "stdout")
        sys.stderr = TextRedirector(self.output_text, "stderr")
    
    def browse_path(self):
        """浏览文件或文件夹"""
        # 尝试打开文件
        path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.png *.jpg *.jpeg *.bmp"), ("所有文件", "*.*")]
        )
        
        # 如果用户没有选择文件，则打开文件夹
        if not path:
            path = filedialog.askdirectory()
        
        if path:
            self.entry_path_var.set(path)
    
    def run_recognition(self):
        """运行识别过程"""
        try:
            # 获取用户输入的参数
            self.collect_parameters()
            
            # 获取源图像路径
            src_img_path = self.entry_path_var.get()
            
            if not src_img_path:
                messagebox.showerror("错误", "请选择图片或文件夹")
                return
            
            # 初始化MIDI映射
            music_name = os.path.basename(src_img_path).split('.')[0]
            midi_map = createMidiInfo(music_name, speed=args["speed"], tonic=args["tonic"])
            
            # 处理单个文件或文件夹
            if os.path.isdir(src_img_path):
                self.process_directory(src_img_path, midi_map)
            else:
                self.process_single_file(src_img_path, midi_map)
            
            # 生成MIDI
            self.generate_midi(music_name, midi_map)
            
            # 生成数据集并打开标注工具
            if self.var_launch_annotation.get():
                self.generate_and_open_annotation_tool(src_img_path)
        
        except Exception as e:
            messagebox.showerror("错误", f"运行过程中发生错误: {str(e)}")
            import traceback
            traceback.print_exc()
    
    def collect_parameters(self):
        """收集用户设置的参数"""
        # 主音和速度
        args["tonic"] = self.entry_tonic_var.get()
        args["speed"] = int(self.entry_speed_var.get())
        
        # 节拍
        args["beat"] = self.entry_beat_var.get()
        
        # 输出中间信息
        args["print_info"] = 1 if self.var_print_info.get() else 0
        
        # 匹配精度
        precision_map = {"高精度 (慢)": 1, "中精度 (推荐)": 2, "低精度 (快)": 3}
        precision_text = self.match_precision_var.get()
        args["match_precision"] = precision_map.get(precision_text, 2)
        args["sampling_rate"] = 1 if args["match_precision"] == 1 else 2
        
        # 线程数
        args["thread_num"] = int(self.entry_thread_num_var.get())
        
        # 字符间距
        args["delta_range"] = float(self.entry_delta_range_var.get())
        
        # 二值化方法
        method_map = {"全局阈值": "global", "OTSU自适应": "OTSU"}
        args["threshold_method"] = method_map.get(self.threshold_method_var.get(), "global")
        
        # 全局阈值
        args["threshold"] = int(self.entry_threshold_var.get())
        
        # 手动调整
        args["wanna_debug"] = self.var_wanna_debug.get()
    
    def process_directory(self, dir_path, midi_map):
        """处理目录中的所有文件"""
        music_sheets = os.listdir(dir_path)
        print(f"路径 '{dir_path}' 下有以下 {len(music_sheets)} 个文件：")
        for num, file_name in enumerate(music_sheets):
            print(f"{num}: {file_name}")
        
        # 获取用户忽略的文件
        ignored_files = self.get_ignored_files(len(music_sheets))
        
        # 处理每个文件
        for num, file_name in enumerate(music_sheets):
            if num in ignored_files:
                continue
            file_path = os.path.join(dir_path, file_name)
            self.process_single_file(file_path, midi_map)
    
    def get_ignored_files(self, total_files):
        """获取用户选择忽略的文件序号"""
        dialog = InputDialog(self, "请输入不需要检测的文件的序号(以空格分隔):", "选择忽略的文件")
        self.wait_window(dialog)
        
        try:
            return set(map(int, dialog.user_input.split()))
        except:
            return set()
    
    def process_single_file(self, file_path, midi_map):
        """处理单个文件"""
        print(f"\n{'='*50}")
        print(f"开始处理: {file_path}")
        
        # 创建识别对象
        self.sheet_recognition = Recognition(
            file_path,
            delta_range=args["delta_range"],
            print_info=args["print_info"],
            match_precision=args["match_precision"],
            thread_num=args["thread_num"],
            sampling_rate=args["sampling_rate"],
            threshold_method=args["threshold_method"],
            threshold=args["threshold"]
        )
        
        # 设置回调函数并运行识别
        self.sheet_recognition.set_callback_function(self.update_output_text)
        self.sheet_recognition.recognize()
        
        # 手动调整
        if args["wanna_debug"]:
            img_name = os.path.basename(file_path).split('.')[0]
            self.sheet_recognition.saveInfo()
            print(f"请查看文件 {img_name}.txt 并订正...")
            if self.wait_for_confirmation("改好后输入Y/y继续:"):
                self.sheet_recognition.readInfo(f"{img_name}.txt")
        
        # 添加到MIDI时间线
        recognition_result = self.sheet_recognition.getTimelineResult()
        midi_map.addTimeline(recognition_result)
        
        print(f"完成处理: {file_path}")
    
    def wait_for_confirmation(self, prompt):
        """等待用户确认"""
        dialog = InputDialog(self, prompt, "确认")
        self.wait_window(dialog)
        return dialog.user_input.upper() == 'Y'
    
    def generate_midi(self, music_name, midi_map):
        """生成MIDI文件"""
        midi_map.mapping()
        midi_timeline = midi_map.getMidiTimeline()
        make_midi = createMidi(1, midi_timeline, midi_name=music_name)
        make_midi.createmidifile()
    
    def generate_and_open_annotation_tool(self, src_img_path):
        """生成数据集并打开标注工具"""
        # 创建专门的数据集存储文件夹
        base_dataset_path = "annotation_dataset"
        if not os.path.exists(base_dataset_path):
            os.makedirs(base_dataset_path)
        
        # 使用图像文件名创建子文件夹
        image_name = os.path.splitext(os.path.basename(src_img_path))[0]
        dataset_path = os.path.join(base_dataset_path, image_name)
        if not os.path.exists(dataset_path):
            os.makedirs(dataset_path)
        
        # 创建切片存储文件夹
        slices_path = os.path.join(dataset_path, "slices")
        if not os.path.exists(slices_path):
            os.makedirs(slices_path)
        
        # 生成数据集
        dataset_gen = DatasetGenerator(
            self.sheet_recognition,
            output_dir=dataset_path,
            slices_dir=slices_path,
            original_image_path=src_img_path
        )
        dataset_gen.generate_dataset()
        
        # 打开标注工具
        self.open_annotation_tool(dataset_path, src_img_path)
    
    def open_annotation_tool(self, dataset_path, original_image_path):
        """打开标注工具"""
        from annotation_tool import AnnotationTool
        
        # 确保原始图像路径有效
        if not os.path.isfile(original_image_path):
            print(f"警告: 原始图像路径无效: {original_image_path}")
            # 尝试从数据集获取
            dataset_file = os.path.join(dataset_path, "dataset.json")
            if os.path.exists(dataset_file):
                try:
                    with open(dataset_file, 'r') as f:
                        dataset = json.load(f)
                    original_image_path = dataset["image_path"]
                except:
                    pass
        
        # 打开标注工具
        tool = AnnotationTool(self, dataset_path, original_image_path)
        self.wait_window(tool)
        
        # 保存纠正后的标注
        tool.save_dataset()
        
        # 计算并显示正确率
        accuracy = tool.calculate_accuracy()
        self.output_text.insert(tk.END, f"\n识别正确率: {accuracy:.2f}%\n", ("stdout",))
        self.output_text.see(tk.END)
    
    def update_output_text(self, line):
        """更新输出文本框"""
        self.output_text.insert(tk.END, line, ("stdout",))
        self.output_text.update()
        self.output_text.see(tk.END)


class InputDialog(tk.Toplevel):
    """输入对话框"""
    def __init__(self, parent, prompt, title):
        super().__init__(parent)
        self.transient(parent)
        self.title(title)
        self.parent = parent
        self.user_input = ""
        
        # 创建界面
        self.prompt = tk.Label(self, text=prompt)
        self.prompt.pack(padx=10, pady=5)
        
        self.user_input_var = tk.StringVar()
        self.user_input_entry = tk.Entry(self, textvariable=self.user_input_var)
        self.user_input_entry.pack(padx=10, pady=5)
        self.user_input_entry.focus_set()
        
        self.ok_button = tk.Button(self, text="确定", command=self.ok)
        self.ok_button.pack(pady=10)
        
        # 居中对话框
        self.update_idletasks()
        x = parent.winfo_rootx() + parent.winfo_width() // 2 - self.winfo_reqwidth() // 2
        y = parent.winfo_rooty() + parent.winfo_height() // 2 - self.winfo_reqheight() // 2
        self.geometry(f"+{x}+{y}")
        
        # 绑定回车键
        self.bind("<Return>", lambda event: self.ok())
    
    def ok(self):
        """确定按钮处理"""
        self.user_input = self.user_input_var.get()
        self.destroy()


class TextRedirector:
    """文本重定向器"""
    def __init__(self, widget, tag="stdout"):
        self.widget = widget
        self.tag = tag
    
    def write(self, str):
        self.widget.insert(tk.END, str, (self.tag,))
        self.widget.see(tk.END)
    
    def flush(self):
        pass


if __name__ == "__main__":
    app = MainApplication()
    app.mainloop()